import cv2
import numpy as np

# 初始化视频捕获对象
cap = cv2.VideoCapture("/media/starkpid/Document/flyList/project/03yumaoqiu/yumaoqiu.mp4")  # 替换为你的视频路径

cv2.namedWindow("Tracking", cv2.WINDOW_NORMAL)  # 创建可调整窗口
cv2.resizeWindow("Tracking", 800, 600)         # 设置初始尺寸为800x600

cv2.namedWindow("Original vs Mask", cv2.WINDOW_NORMAL)  # 创建可调整窗口
cv2.resizeWindow("Original vs Mask", 800, 1200)         # 设置初始尺寸为800x600

# 获取视频帧率
fps = cap.get(cv2.CAP_PROP_FPS)
print("视频帧率:", fps, "FPS")

# 白色物体HSV阈值范围（H: 0-180, S: 0-255, V: 0-255）
lower_white = np.array([0, 0, 200])
upper_white = np.array([255, 30, 255])

# 多目标跟踪相关变量
tracked_objects = []  # 存储跟踪对象信息（ID + 位置）
current_id = 0        # 对象ID计数器
MAX_DISTANCE = 50     # 最大关联距离（像素）

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    # 转换为HSV颜色空间并创建掩膜
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    mask = cv2.inRange(hsv, lower_white, upper_white)

    # 形态学操作（去噪）
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=2)

    # 查找轮廓
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 将mask转换为3通道用于并排显示
    mask_bgr = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
    combined = np.hstack((frame, mask_bgr))
    cv2.imshow('Original vs Mask', combined)
    
    # 提取当前帧的物体中心点
    current_centers = []
    for cnt in contours:
        if cv2.contourArea(cnt) < 500:  # 过滤小面积噪声
            continue
        
        # 计算中心点
        M = cv2.moments(cnt)
        if M["m00"] != 0:
            cx = int(M["m10"] / M["m00"])
            cy = int(M["m01"] / M["m00"])
            current_centers.append((cx, cy))
            
            # 绘制包围盒
            x, y, w, h = cv2.boundingRect(cnt)
            cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0), 2)

    # 多目标跟踪：数据关联
    used_ids = []
    updated_objects = []
    
    for center in current_centers:
        # 寻找最近已知对象
        min_dist = MAX_DISTANCE
        matched_obj = None
        
        for obj in tracked_objects:
            dx = center[0] - obj["position"][0]
            dy = center[1] - obj["position"][1]
            distance = np.sqrt(dx**2 + dy**2)
            
            if distance < min_dist:
                min_dist = distance
                matched_obj = obj
        
        # 更新或创建对象
        if matched_obj:
            updated_objects.append({
                "id": matched_obj["id"],
                "position": center
            })
            used_ids.append(matched_obj["id"])
        else:
            updated_objects.append({
                "id": current_id,
                "position": center
            })
            used_ids.append(current_id)
            current_id += 1
    
    # 移除丢失的对象
    tracked_objects = updated_objects

    # 绘制跟踪信息
    for obj in tracked_objects:
        cx, cy = obj["position"]
        cv2.circle(frame, (cx, cy), 5, (0, 0, 255), -1)
        cv2.putText(frame, f"ID:{obj['id']}", 
                   (cx - 30, cy - 20), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, 
                   (0, 0, 255), 2)

    # 显示结果
    cv2.imshow("Tracking", frame)
    if cv2.waitKey(0) & 0xFF == ord("q"):
        break

cap.release()
cv2.destroyAllWindows()